#include "BkDrv.h"
#include "bklib.h"
//////////////////////////////////////////////////////////////////////////
// Following structures should be allocated globally and should be pre-initialized with zeroes 
static	KLDR_DATA_TABLE_ENTRY	MyLdrDataTableEntry = { 0 };
static	BOOT_DRIVER_LIST_ENTRY	MyBootDriverListEntry = { 0 };
static  UNICODE_STRING			MyRegistryPath = { sizeof(wczRegPath) - sizeof(WCHAR), sizeof(wczRegPath), wczRegPath };

// Global variables
PVOID	BkNtKernelBase = NULL;	// NT kernel base
PVOID	BkImageBase = NULL;	// Current driver base
#define  MAX_DLL_NAME	0x40	// bytes
//	Pointer to the list of loded modules. Used by PeSupGetModuleBase function. Must be initialized by the caller.
PLIST_ENTRY	PeSupLoadedModuleList = NULL;

//////////////////////////////////////////////////////////////////////////
static NTSTATUS ImportScanLoop(
	PCHAR	ModuleBase,
	ULONG	SizeOfImage,
	PVOID	ImportedModuleBase,
	ULONG   rvaINT,
	ULONG   rvaIAT
)
{
	NTSTATUS	ntStatus = STATUS_SUCCESS;
	PIAT_ENTRY	pIatEntry = NULL;
	PIMAGE_IMPORT_BY_NAME   pOrdinalName;
	PIMAGE_THUNK_DATA		pINT;
	PIMAGE_THUNK_DATA		pIAT;
	IAT_ENTRY				IatEntry;

	if (rvaINT == 0)   // No Characteristics field?
	{
		// Yes! Gotta have a non-zero FirstThunk field then.
		rvaINT = rvaIAT;

		if (rvaINT == 0)   // No FirstThunk field?  Ooops!!!
			return(STATUS_INVALID_IMAGE_FORMAT);
	}

	// Adjust the pointer to point where the tables are in the
	// mem mapped file.
	pINT = (PIMAGE_THUNK_DATA)PeSupRvaToVa(rvaINT, ModuleBase);
	if (!pINT)
		return(STATUS_INVALID_IMAGE_FORMAT);

	pIAT = (PIMAGE_THUNK_DATA)PeSupRvaToVa(rvaIAT, ModuleBase);

	while (TRUE) // Loop forever (or until we break out)
	{
		if (pINT->u1.AddressOfData == 0)
			break;

		if (IMAGE_SNAP_BY_ORDINAL(pINT->u1.Ordinal) == FALSE)
		{
			pOrdinalName = (PIMAGE_IMPORT_BY_NAME)PeSupRvaToVa((ULONG)pINT->u1.AddressOfData, ModuleBase);
		}
		else if (pINT->u1.Ordinal >= (ULONG_PTR)ModuleBase &&
			pINT->u1.Ordinal < ((ULONG_PTR)ModuleBase + SizeOfImage))
		{
			pOrdinalName = (PIMAGE_IMPORT_BY_NAME)((ULONG_PTR)pINT->u1.AddressOfData);
		}
		else
			pOrdinalName = NULL;


		if (pOrdinalName)
		{
			IatEntry = (IAT_ENTRY)PeSupGetFunctionAddress(ImportedModuleBase, (PCHAR)&pOrdinalName->Name);
			if (IatEntry)
			{
				pIatEntry = &pIAT->u1.Function;
				*pIatEntry = IatEntry;
			}
			else
			{
				ntStatus = STATUS_PROCEDURE_NOT_FOUND;
				break;
			}
		}	// if (pOrdinalName)

		pINT++;         // advance to next thunk
		pIAT++;         // advance to next thunk
	} // while (TRUE)	

	return(ntStatus);
}


//
//	Resolves specified image import based on the specified list of loaded modules.
//
NTSTATUS	PeSupResolveImageImport(
	PVOID	ImageBase,
	PPE_GET_MODULE_BASE	GetModuleBase
)
{
	ULONG   importsStartRVA;
	PCHAR   ModuleName;
	PCHAR	ModuleBase = (PCHAR)ImageBase;
	PVOID	ImportedModuleBase = NULL;
	NTSTATUS	ntStatus = STATUS_SUCCESS;
	//PIAT_ENTRY	pIatEntry = NULL;
	PIMAGE_IMPORT_DESCRIPTOR pImportDesc;
	//UNREFERENCED_PARAMETER(pIatEntry);

	PIMAGE_NT_HEADERS	PEHeader = (PIMAGE_NT_HEADERS)PeSupGetImagePeHeader(ModuleBase);

	// Get the import table RVA from the data dir
	importsStartRVA =
		PEHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;

	if (!importsStartRVA)
		return(STATUS_INVALID_PARAMETER);

	pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR)PeSupRvaToVa(importsStartRVA, ModuleBase);

	if (!pImportDesc)
		return(STATUS_INVALID_PARAMETER);

	// Find the import descriptor containing references to callee's functions
	for (; pImportDesc->Name; pImportDesc++)
	{
		ModuleName = (PCHAR)PeSupRvaToVa(pImportDesc->Name, ModuleBase);
		if (ModuleName)
		{
			if (GetModuleBase == NULL)
				ImportedModuleBase = PeSupGetModuleBase(ModuleName);
			else
				ImportedModuleBase = (GetModuleBase)(ModuleName);

			if (ImportedModuleBase)
			{
				ntStatus = ImportScanLoop(ModuleBase, PEHeader->OptionalHeader.SizeOfImage, ImportedModuleBase,
					pImportDesc->OriginalFirstThunk, pImportDesc->FirstThunk);

				if (!(NT_SUCCESS(ntStatus)))
					break;
			}	// if (ImportedModuleBase)
			else
			{
				ntStatus = STATUS_INVALID_IMPORT_OF_NON_DLL;
				break;

			}
		} // if(ModuleName = (PCHAR)PeSupRvaToVa
	} // for (; pImportDesc->Name; pImportDesc++) 

	return(STATUS_SUCCESS);
}


//
//	Searches for the export table entry of the specified function within the specified module. 
//	Returns pointer to the found entry or NULL if there is no such function within module's export.
//
PEXPORT_ENTRY	PeSupGetExportEntry(
	PVOID	TargetModule,		// Image base of the target module (where exported function located)
	PCHAR   ExportedFunction	// exported function name
)
{
	PIMAGE_DATA_DIRECTORY	pExportData;
	PIMAGE_EXPORT_DIRECTORY pExportDir;
	PCHAR	ModuleBase = (PCHAR)TargetModule;
	PULONG  ppFunctions, ppNames;
	PUSHORT pOrdinals;
	ULONG	i, NumberOfNames;
	PEXPORT_ENTRY	pFoundEntry = NULL;


	pExportData = PeSupGetDirectoryEntryPtr((PIMAGE_NT_HEADERS)(ModuleBase + ((PIMAGE_DOS_HEADER)TargetModule)->e_lfanew), IMAGE_DIRECTORY_ENTRY_EXPORT);
	pExportDir = (PIMAGE_EXPORT_DIRECTORY)(ModuleBase + pExportData->VirtualAddress);

	if (pExportData->VirtualAddress == 0 || pExportData->Size == 0 || pExportDir->AddressOfFunctions == 0 || pExportDir->AddressOfNames == 0)
		return NULL;

	// Get names, functions and ordinals arrays pointers
	ppFunctions = (PULONG)(ModuleBase + (ULONG)pExportDir->AddressOfFunctions);
	ppNames = (PULONG)(ModuleBase + (ULONG)pExportDir->AddressOfNames);
	pOrdinals = (PUSHORT)(ModuleBase + (ULONG)pExportDir->AddressOfNameOrdinals);

	NumberOfNames = pExportDir->NumberOfNames;

	// Walk the export table entries
	for (i = 0; i < NumberOfNames; ++i)
	{
		// Check if function name matches current entry
		if (!__strcmp(ModuleBase + *ppNames, ExportedFunction))
		{
			pFoundEntry = (PEXPORT_ENTRY)&ppFunctions[*pOrdinals];
			break;
		}
		ppNames++;
		pOrdinals++;
	}

	return(pFoundEntry);
}


//				
// Retrieves the address of an exported function or variable from the specified module.
//
PVOID	PeSupGetFunctionAddress(
	PVOID	ImageBase,			// target module base
	PCHAR	FunctionName		// target function name
)
{
	PVOID			Address = NULL;
	PEXPORT_ENTRY	pExportEntry = { 0 };
	PIMAGE_DATA_DIRECTORY pExportData;

	pExportData = PeSupGetDirectoryEntryPtr((PIMAGE_NT_HEADERS)((PCHAR)ImageBase + ((PIMAGE_DOS_HEADER)ImageBase)->e_lfanew), IMAGE_DIRECTORY_ENTRY_EXPORT);

	pExportEntry = PeSupGetExportEntry(ImageBase, FunctionName);
	if (pExportEntry)
	{
		Address = (PVOID)((PCHAR)ImageBase + *pExportEntry);

	}	// if (pExportEntry = PeSupGetExportEntry(ImageBase, FunctionName))

	return(Address);
}


//
// Retrieves a module base for the specified module from the specified list of loaded modules.
//
PVOID	PeSupGetModuleBase(
	PCHAR	ModuleName
)
{
	BOOL	Found = FALSE;
	CHAR	ShortName[MAX_DLL_NAME];
	PVOID	ModuleBase = NULL;
	PLIST_ENTRY	pEntry;
	PKLDR_DATA_TABLE_ENTRY	LdrEntry = NULL;
	ULONG	NameLen = (ULONG)__strlen(ModuleName);

	pEntry = PeSupLoadedModuleList->Flink;

	while (pEntry != PeSupLoadedModuleList)
	{
		LdrEntry = CONTAINING_RECORD(pEntry, KLDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
		if (LdrEntry->BaseImageName.Length < MAX_DLL_NAME)
		{
			ULONG i;
			PCHAR Dot;

			for (i = 0; i < LdrEntry->BaseImageName.Length; i++)
				ShortName[i] = (CHAR)LdrEntry->BaseImageName.Buffer[i];
			ShortName[NameLen] = 0;

			if (__stricmp(ShortName, ModuleName) == 0)
			{
				Found = TRUE;
				break;
			}
			Dot = __strchr(ShortName, '.');
			if (Dot)
			{
				Dot[0] = 0;
				if (__stricmp(ShortName, ModuleName) == 0)
				{
					Found = TRUE;
					break;
				}
			}	// if (Dot = strchr(ShortName, '.'))
		}	// if (LdrEntry->BaseDllName.Length < MAX_DLL_NAME)
		pEntry = pEntry->Flink;
	}	//  while(pEntry != PeSupLoadedModuleList)

	if (Found)
		ModuleBase = LdrEntry->ImageBase;

	return(ModuleBase);
}
//////////////////////////////////////////////////////////////////////////
//
// Retrieves the address of an exported function or variable from the specified module.
//
PVOID	BkGetFunctionAddress(
	PVOID		ImageBase,			// target module base
	PCHAR		FunctionName		// target function name
)
{
	return(PeSupGetFunctionAddress(ImageBase, FunctionName));
}


//
// Retrieves a module base for the specified module.
//
PVOID	BkGetModuleBase(
	PCHAR	ModuleName
)
{
	return(PeSupGetModuleBase(ModuleName));
}


//
//	Returns the export entry address for the specified function of the specified module.
//
PVOID	BkGetExportEntry(
	PVOID	ImageBase,
	PCHAR	FunctionName
)
{
	return(PeSupGetExportEntry(ImageBase, FunctionName));
}


//
//	Resolves specified executable image import.
//
NTSTATUS BkResolveImageImport(
	PVOID	ImageBase,
	PBK_GET_MODULE_BASE	pBkGetModuleBase
)
{
	return(PeSupResolveImageImport(ImageBase, pBkGetModuleBase));
}


/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//	This functon used instead of memcpy() because original memcpy() for x64 compiled with SSE instuctions.
//
static void* xmemcpy(
	void * dst,
	const void * src,
	size_t count
)
{
	void * ret = dst;

	/*
	* copy from lower addresses to higher addresses
	*/
	while (count--) {
		*(char *)dst = *(char *)src;
		dst = (char *)dst + 1;
		src = (char *)src + 1;
	}
	return(ret);
}

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//	This functon used instead of strcpy() because original strcpy() for x64 compiled with SSE instuctions.
//
static int xstrcpy(
	char * dst,
	const char * src
)
{
	int	count = 0;
	char c;

	while (c = src[0])
	{
		dst[0] = c;
		count += 1;
		dst += 1;
		src += 1;
	}
	dst[0] = c;
	return(count);
}
//////////////////////////////////////////////////////////////////////////
//
//	BK initialization routine. Normally called twice: 
//	 - when BK loads the driver, is this case pDriverObject == 0 and pusRegistryPath is pointer to PBK_LDR_CONTEXT structure.
//	 - when the driver being registered by the OS.
//
NTSTATUS BkInitialize(
	PVOID					pDriverObject,
	PUNICODE_STRING			pusRegistryPath,
	PBK_INSTALL_CALLBACK	pfnInstallRoutine,
	PBK_DRIVER_CALLBACK		pfnDriverRoutine
)
{
	NTSTATUS	ntStatus = STATUS_INVALID_PARAMETER;

	if (pDriverObject)
	{
		// Driver being registered by the OS, calling driver startup routine
		if (pfnDriverRoutine)
			ntStatus = (pfnDriverRoutine)(pDriverObject, pusRegistryPath);
	}
	else
	{
		// pDriverObject is NULL means we are called by the BK and 
		//	pusRegistryPath parameter contains pointer to BK_LDR_CONTEXT structure.
		if (pfnInstallRoutine)
			ntStatus = (pfnInstallRoutine)(pusRegistryPath);
	}

	return(ntStatus);
}

NTSTATUS	BkInstall(
	PBK_LDR_CONTEXT	LdrContext
)
{
	NTSTATUS	ntStatus = STATUS_SUCCESS;
	PBOOT_DRIVER_LIST_ENTRY	pBootDriverListEntry;
	PKLDR_DATA_TABLE_ENTRY	pLdrDataTableEntry;
	IMAGE_DATA_DIRECTORY *ExceptionTable;
	ULONG	Length;

	BkNtKernelBase = LdrContext->NtImageBase;
	BkImageBase = LdrContext->DriverBase;

	// Initializing loaded module list pointer
	PeSupLoadedModuleList = &LdrContext->NtLoaderBlock->LoadOrderListHead;

	// Resolving current driver's import
	PeSupResolveImageImport(BkImageBase, NULL);

	pBootDriverListEntry = CONTAINING_RECORD(LdrContext->NtLoaderBlock->BootDriverListHead.Flink, BOOT_DRIVER_LIST_ENTRY, Link);
	pLdrDataTableEntry = CONTAINING_RECORD(LdrContext->NtLoaderBlock->LoadOrderListHead.Flink, KLDR_DATA_TABLE_ENTRY, InLoadOrderLinks);

	// Initializing our driver boot list entry
	// Copying whole structure from any other valid driver. Actually we need any driver path only.
	xmemcpy(&MyBootDriverListEntry, pBootDriverListEntry, sizeof(BOOT_DRIVER_LIST_ENTRY));
	//	Initializing list entry and pointer
	InitializeListHead(&MyBootDriverListEntry.Link);
	MyBootDriverListEntry.LdrEntry = (PKLDR_DATA_TABLE_ENTRY)&MyLdrDataTableEntry;
	// IopInitializeBootDrivers requires driver registry path to be valid and accessable.
	xmemcpy(&MyBootDriverListEntry.RegistryPath, &MyRegistryPath, sizeof(UNICODE_STRING));

	// Initializing our driver loader data table entry
	Length = sizeof(KLDR_DATA_TABLE_ENTRY);
	// Copying whole structure from any other valid driver. We need valid flags and system pointers here.
	xmemcpy(&MyLdrDataTableEntry, pLdrDataTableEntry, Length);
	InitializeListHead(&MyLdrDataTableEntry.InLoadOrderLinks);

	// Following fields have the same offsets on all OS version
	MyLdrDataTableEntry.ImageBase = LdrContext->DriverBase;
	MyLdrDataTableEntry.EntryPoint = LdrContext->DriverEntry;
	MyLdrDataTableEntry.SizeOfImage = (ULONG)LdrContext->DriverSize;
	//SEH
	ExceptionTable = PeSupGetDirectoryEntryPtr(LdrContext->DriverBase, IMAGE_DIRECTORY_ENTRY_EXCEPTION);
	if (ExceptionTable)
	{
		MyLdrDataTableEntry.ExceptionTable = PeSupRvaToVa(ExceptionTable->VirtualAddress, LdrContext->DriverBase);
		MyLdrDataTableEntry.ExceptionTableSize = ExceptionTable->Size;
	}
	// Registering our driver within the loader parameter block
	InsertTailList(&LdrContext->NtLoaderBlock->LoadOrderListHead, &MyLdrDataTableEntry.InLoadOrderLinks);
	InsertTailList(&LdrContext->NtLoaderBlock->BootDriverListHead, &MyBootDriverListEntry.Link);

	return(ntStatus);
}
//////////////////////////////////////////////////////////////////////////
NTSTATUS	BkStartThread(PKSTART_ROUTINE TrdRoutine)
{
	HANDLE		hThread;
	NTSTATUS	ntStatus;
	OBJECT_ATTRIBUTES	oa = { 0 };

	InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
	ntStatus = PsCreateSystemThread(&hThread, GENERIC_ALL, &oa, NULL, NULL, TrdRoutine, NULL);
	if (NT_SUCCESS(ntStatus))
		ZwClose(hThread);
	return(ntStatus);
}
//////////////////////////////////////////////////////////////////////////
/***
*strcmp - compare two strings, returning less than, equal to, or greater than
*
*Purpose:
*       STRCMP compares two strings and returns an integer
*       to indicate whether the first is less than the second, the two are
*       equal, or whether the first is greater than the second.
*
*       Comparison is done byte by byte on an UNSIGNED basis, which is to
*       say that Null (0) is less than any other character (1-255).
*
*Entry:
*       const char * src - string for left-hand side of comparison
*       const char * dst - string for right-hand side of comparison
*
*Exit:
*       returns -1 if src <  dst
*       returns  0 if src == dst
*       returns +1 if src >  dst
*
*Exceptions:
*
*******************************************************************************/

int __cdecl __strcmp(
	const char * src,
	const char * dst
)
{
	int ret = 0;
	ret = *(unsigned char *)src - *(unsigned char *)dst;
	while (!ret && *dst)
	{
		++src, ++dst;
		ret = *(unsigned char *)src - *(unsigned char *)dst;
	}
	if (ret < 0)
		ret = -1;
	else if (ret > 0)
		ret = 1;

	return(ret);
}


/***
*wcscmp - compare two wchar_t strings,
*        returning less than, equal to, or greater than
*
*Purpose:
*       wcscmp compares two wide-character strings and returns an integer
*       to indicate whether the first is less than the second, the two are
*       equal, or whether the first is greater than the second.
*
*       Comparison is done wchar_t by wchar_t on an UNSIGNED basis, which is to
*       say that Null wchar_t(0) is less than any other character.
*
*Entry:
*       const wchar_t * src - string for left-hand side of comparison
*       const wchar_t * dst - string for right-hand side of comparison
*
*Exit:
*       returns -1 if src <  dst
*       returns  0 if src == dst
*       returns +1 if src >  dst
*
*Exceptions:
*
*******************************************************************************/

int __cdecl __wcscmp(
	const wchar_t * src,
	const wchar_t * dst
)
{
	int ret = 0;
	ret = (int)(*src - *dst);
	while (!ret && *dst) {
		++src, ++dst;
		ret = (int)(*src - *dst);
	}

	if (ret < 0)
		ret = -1;
	else if (ret > 0)
		ret = 1;

	return(ret);
}



int __cdecl __stricmp(
	const char * dst,
	const char * src
)
{
	int f, l;

	do
	{
		if (((f = (unsigned char)(*(dst++))) >= 'A') && (f <= 'Z'))
			f -= 'A' - 'a';
		if (((l = (unsigned char)(*(src++))) >= 'A') && (l <= 'Z'))
			l -= 'A' - 'a';
	} while (f && (f == l));

	return(f - l);
}

int __cdecl __wcsicmp(
	const wchar_t * dst,
	const wchar_t * src
)
{
	wchar_t f, l;

	do {
		f = __ascii_towlower(*dst);
		l = __ascii_towlower(*src);
		dst++;
		src++;
	} while ((f) && (f == l));

	return (int)(f - l);
}


/***
*char *strchr(string, c) - search a string for a character
*
*Purpose:
*       Searches a string for a given character, which may be the
*       null character '\0'.
*
*Entry:
*       char *string - string to search in
*       char c - character to search for
*
*Exit:
*       returns pointer to the first occurence of c in string
*       returns NULL if c does not occur in string
*
*Exceptions:
*
*******************************************************************************/

char * __cdecl __strchr(
	const char * string,
	int ch
)
{
	while (*string && *string != (char)ch)
		string++;

	if (*string == (char)ch)
		return((char *)string);
	return(NULL);
}



/***
*wchar_t *wcschr(string, c) - search a string for a wchar_t character
*
*Purpose:
*       Searches a wchar_t string for a given wchar_t character,
*       which may be the null character L'\0'.
*
*Entry:
*       wchar_t *string - wchar_t string to search in
*       wchar_t c - wchar_t character to search for
*
*Exit:
*       returns pointer to the first occurence of c in string
*       returns NULL if c does not occur in string
*
*Exceptions:
*
*******************************************************************************/

wchar_t * __cdecl __wcschr(
	const wchar_t * string,
	wchar_t ch
)
{
	while (*string && *string != (wchar_t)ch)
		string++;

	if (*string == (wchar_t)ch)
		return((wchar_t *)string);
	return(NULL);
}



/***
*strlen - return the length of a null-terminated string
*
*Purpose:
*       Finds the length in bytes of the given string, not including
*       the final null character.
*
*Entry:
*       const char * str - string whose length is to be computed
*
*Exit:
*       length of the string "str", exclusive of the final null byte
*
*Exceptions:
*
*******************************************************************************/

size_t __cdecl __strlen(
	const char * str
)
{
	const char *eos = str;

	while (*eos++);

	return(eos - str - 1);
}
//////////////////////////////////////////////////////////////////////////